import os
import numpy as np
import torchaudio
import torch
import torch.nn as nn
import torch.nn.functional as F
import joblib
from sklearn.svm import LinearSVC
from sklearn.metrics import accuracy_score
from torch.utils.data import Dataset, DataLoader

class SoundNetModel(nn.Module):
    def __init__(self, net_path):
        super(SoundNetModel, self).__init__()
        self.net = torch.load(net_path)
        self.net = nn.Sequential(*list(self.net.children())[:-4])
        
    def forward(self, x):
        return self.net(x)

class ESC50Dataset(Dataset):
    def __init__(self, file_list, data_root, sound_length):
        self.file_list = file_list
        self.data_root = data_root
        self.sound_length = sound_length
        
        with open(file_list, 'r') as f:
            self.files = f.readlines()
        
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        line = self.files[idx].strip().split()
        snd_path = os.path.join(self.data_root, line[0])
        label = int(line[1])
        snd, sr = torchaudio.load(snd_path)
        snd = snd.mul(2**-22)
        
        # Circular repeat
        rep = int(np.ceil(((self.sound_length + 1) * sr) / snd.size(1)))
        snd = snd.repeat(rep, 1)
        
        # Cut at the desired length
        if snd.size(1) > self.sound_length * sr:
            snd = snd[:, :self.sound_length * sr]
        
        return snd, label

def read_dataset(file, data_root, sound_length):
    dataset = ESC50Dataset(file, data_root, sound_length)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
    features, labels = [], []
    
    for snd, label in dataloader:
        snd = snd.view(1, 1, -1, 1).float().cuda()
        with torch.no_grad():
            feat = net(snd)
        feat = feat.cpu().squeeze()
        features.append(feat)
        labels.append(label.item())
    
    return torch.stack(features), torch.tensor(labels)

def convert_to_liblinear(features, labels):
    ds = []
    for i in range(features.size(0)):
        label = labels[i].item()
        feat = features[i].view(-1)
        nz = torch.nonzero(feat > 0).squeeze()
        
        if nz.dim() == 0:
            nz = torch.tensor([0])
            val = torch.tensor([0.0])
        else:
            val = feat[nz]
        
        ds.append((label, (nz.int().numpy(), val.float().numpy())))
    
    return ds

def main(opt):
    global net
    net = SoundNetModel(opt.net)
    net.cuda()
    net.eval()
    
    print('extracting layer', opt.layer)
    
    train_data = read_dataset(os.path.join(opt.data_root, 'splits', 'train' + str(opt.fold) + '.txt'), opt.data_root, opt.sound_length)
    test_data = read_dataset(os.path.join(opt.data_root, 'splits', 'test' + str(opt.fold) + '.txt'), opt.data_root, opt.sound_length)
    
    train_data_ll = convert_to_liblinear(*train_data)
    test_data_ll = convert_to_liblinear(*test_data)
    
    print('train')
    model = LinearSVC(C=0.01)
    X_train, y_train = zip(*train_data_ll)
    X_train = [np.concatenate((nz, val)) for nz, val in X_train]
    model.fit(X_train, y_train)
    
    print('training accuracy:')
    train_preds = model.predict(X_train)
    train_accuracy = accuracy_score(y_train, train_preds)
    print(train_accuracy)
    
    print('testing accuracy:')
    X_test, y_test = zip(*test_data_ll)
    X_test = [np.concatenate((nz, val)) for nz, val in X_test]
    test_preds = model.predict(X_test)
    test_accuracy = accuracy_score(y_test, test_preds)
    print(test_accuracy)
    
    with open(opt.outfile, "a") as f:
        f.write(f"Net: {opt.net}\t Layer: {opt.layer}\t TimeWindow: {opt.timewindow}\t C: 0.01\t Fold: {opt.fold}\t Accuracy: {test_accuracy}\n")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Evaluate Soundnet on ESC50')
    parser.add_argument('-net', default='models/soundnet8_final.pth', type=str, help='net name')
    parser.add_argument('-layer', default=18, type=int, help='layer')
    parser.add_argument('-timewindow', default=1, type=int, help='initial time window')
    parser.add_argument('-fold', default=1, type=int, help='esc50 fold')
    parser.add_argument('-outfile', default='results_esc50.txt', type=str, help='output filename')
    opt = parser.parse_args()
    
    main(opt)
